package com.gnerv.pabu.basic.agent;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.NotFoundException;
import javassist.bytecode.CodeAttribute;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.reflect.Method;
import java.security.ProtectionDomain;

public class Transformer implements ClassFileTransformer {
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        if (className == null) {
            //返回null，将会使用原生class。
            return null;
        }
        if (className.startsWith("java") ||
                className.startsWith("javax") ||
                className.startsWith("jdk") ||
                className.startsWith("sun") ||
                className.startsWith("com/sun") ||
                className.startsWith("com/intellij") ||
                className.startsWith("org/jetbrains") ||
                className.startsWith("com/thoreauz/agent")
        ) {
            // 不对JDK类以及agent类增强
            return null;
        }
        try {
            if(className.startsWith("com/gnerv/pabu/basic/core/controller")){
                CtClass controller = ClassPool.getDefault().get(className.replace("/", "."));
                boolean hasAnnotation = controller.hasAnnotation("org.springframework.web.bind.annotation.RestController");
                if (hasAnnotation) {
                    CtMethod[] methods = controller.getMethods();
                    for (CtMethod method : methods) {
                        CodeAttribute ca = method.getMethodInfo2().getCodeAttribute();
                        if (ca == null) {
                            continue;
                        }
                        if (!method.isEmpty()) {
                            method.insertBefore("System.out.println(\"开始 : " + method.getName() + "\");");
                            method.insertAfter("System.out.println(\"结束 : " + method.getName() + "\");");

                        }
                    }
                    return controller.toBytecode();
                }
            }
        } catch (NotFoundException | CannotCompileException | IOException e) {
            e.printStackTrace();
        }
        return classfileBuffer;
    }
}